/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

/*! \file multicoloring.h
 *  \brief Multicoloring of a matrix.
 */

#include <iostream>
#include <set>
#include <cusp/coo_matrix.h>

#pragma once

namespace cusp {
namespace graph {
namespace detail {
namespace host {

struct adjacency_matrix
{
    std::vector<int> degrees;
    std::vector<int> neighbours;

    adjacency_matrix() {}
    inline size_t num_vertices( ) const { return degrees.size( ); }
    inline void resize( int n, int num_neighbours ) 
    {
        degrees.resize( n );
        neighbours.resize( num_neighbours );
    }
};

struct edge
{
    int u, v;
    edge( int uu, int vv ) : u( uu ), v( vv ) {}
    bool operator < ( const edge &e ) const { return u < e.u || ( u == e.u && v < e.v ); }
};

template< typename Array >
bool dbg_coloring_is_valid( const adjacency_matrix &g, const Array &coloring )
{
    for( int u = 0, i = 0 ; u < g.num_vertices( ) ; ++u )
        for( int k = 0 ; k < g.degrees[u] ; ++k, ++i )
            if( coloring[u] == coloring[g.neighbours[i]] )
                return false;
    return true;
}

template< typename Matrix >
void build_adjacency_matrix( const Matrix &a, adjacency_matrix &g, cusp::coo_format )
{
    typedef std::set<edge> edge_set;
    typedef edge_set::iterator edge_set_iterator;
    edge_set edges;
    size_t order = std::max( a.num_rows, a.num_cols );
    for( int i = 0 ; i != a.num_entries ; ++i ) 
    {
        int u = a.row_indices[i], v = a.column_indices[i];
        if( u == v )
            continue;
        edges.insert( edge( u, v ) );
        edges.insert( edge( v, u ) );
    }
    g.resize( order, edges.size( ) );
    int u = 0, i = 0, degree = 0;
    for( edge_set_iterator it = edges.begin( ) ; it != edges.end( ) ; ++it )
    {
        if( it->u != u ) 
        {
            g.degrees[u++] = degree;
            degree = 0;
            assert( u == it->u );
        }
        g.neighbours[i++] = it->v;
        degree++;
    }
    g.degrees[u] = degree;
}

template< typename Matrix, typename Vector, typename Format >
void greedy_multicoloring( const Matrix &a, Vector &coloring, Format format )
{
    adjacency_matrix g;
    build_adjacency_matrix( a, g, format );
    std::fill( coloring.begin( ), coloring.end( ), std::numeric_limits<int>::max( ) );

    // For each vertex, find the color s.t. c = min{ k > 0 : color[v] != k, where (u,v) in E(A) }.
    coloring[0] = 0;
    int i = g.degrees[0];
    for( int u = 1, end = g.num_vertices( ) ; u < end ; ++u )
    {
        typedef std::set<int> color_set;
        typedef color_set::iterator color_set_it;
        color_set colors;
        bool done = false;
        for( int k = 0 ; !done && k < g.degrees[u] ; ++k )
        {
            assert( i+k < g.neighbours.size( ) );
            int v = g.neighbours[i+k];
            if( v < u )
                colors.insert( coloring[v] );
            done = v >= u;
        }
        i += g.degrees[u];

        // The set of colors is already sorted.
        int color = 0;
        color_set_it it = colors.begin( );
        for( ; it != colors.end( ) && *it == color ; ++it, ++color )
            ;
        assert( u < coloring.size( ) );
        coloring[u] = color;
    }
    assert( dbg_coloring_is_valid( g, coloring ) );
}

// Ordering based on the coloring.
template< typename Array1, typename Array2 >
struct color_based_order
{
    typedef typename Array2::value_type value_type;
    const Array1 &c;
    color_based_order( const Array1 &col ) : c( col ) {}
    inline bool operator( )( const value_type &u, const value_type &v ) const { return c[u] < c[v]; }
};

// Build the permutation from the coloring.
template< typename Array1, typename Array2 >
void build_permutation_from_coloring( const Array1 &coloring, Array2 &p, cusp::host_memory )
{
    thrust::host_vector<typename Array2::value_type> tmp( p.size( ) );
    thrust::sequence( tmp.begin( ), tmp.end( ) );
    thrust::sort( tmp.begin( ), tmp.end( ), color_based_order<Array1, Array2>( coloring ) );
    for( int i = 0, end = tmp.size( ) ; i < end ; ++i )
        p[tmp[i]] = i;
}

} // namespace host

// Run the computation on the host.
template< typename Matrix, typename Array1d >
void greedy_multicoloring( const Matrix &a, Array1d &coloring, cusp::host_memory )
{
    host::greedy_multicoloring( a, coloring, typename Matrix::format( ) );
}

// Build the permutation from the coloring.
template< typename Array1, typename Array2 >
void build_permutation_from_coloring( const Array1 &coloring, Array2 &p, cusp::host_memory )
{
    host::build_permutation_from_coloring( coloring, p );
}

// The device path redirects the computation to the host (for the moment).
template< typename Matrix, typename Array1d >
void greedy_multicoloring( const Matrix &a, Array1d &coloring, cusp::device_memory )
{
    typedef typename Matrix::template rebind<cusp::host_memory>::type host_matrix;
    typedef cusp::array1d<typename Array1d::value_type, cusp::host_memory> host_array1d;
    host_matrix h_a;
    cusp::copy( a, h_a );
    host_array1d h_coloring;
    host::greedy_multicoloring( h_a, h_coloring, typename Matrix::format( ) );
    cusp::copy( h_coloring, coloring );
    cusp::copy( h_a, a );
}

// The device path redirects the computation to the host (for the moment).
template< typename Array1, typename Array2 >
void build_permutation_from_coloring( const Array1 &coloring, Array2 &p, cusp::device_memory )
{
    assert( false );
}

} // namespace detail

template< typename Matrix, typename Array1d >
void greedy_multicoloring( const Matrix &a, Array1d &coloring )
{
    detail::greedy_multicoloring( a, coloring, typename Matrix::memory_space( ) );
}

template< typename Array1, typename Array2 >
void build_permutation_from_coloring( const Array1 &coloring, Array2 &p )
{
    detail::build_permutation_from_coloring( coloring, p, typename Array1::memory_space( ) );
}

} // namespace graph
} // namespace cusp

